using System;
using System.Collections.Generic;
using System.Text;
using System.Collections;
using System.Drawing;
using System.Diagnostics;
using System.Drawing.Imaging;

namespace ImageManipulation
{  
    /// <summary>
    /// Quantize using median cut.
    /// </summary>
    public unsafe class MedianCutQuantizer : Quantizer
    {
        /// <summary>
        /// Construct the median cut quantizer
        /// </summary>
        /// <remarks>
        /// The median cut algorithm  is slightly different from the two pass algorithm generalization
        /// made by Quantizer super class.  In the second pass, instead of iterate through each pixel
        /// of source image in order and convert them from color to indices, the clusters and pixels
        /// within those clusters are iterated.
        /// </remarks>
        /// <param name="maxColors">The maximum number of colors to return</param>
        public MedianCutQuantizer(int maxColors)
            : base(false)
        {
            // GDI+ support at most 2^8 colors in indexed pixel format
            if (maxColors > 256)
                throw new ArgumentOutOfRangeException("maxColors", maxColors, "The number of colors should be less than 256");

            _maxColors = maxColors;

            _splittables = new Queue<Cluster>();
            _nonsplittables = new List<Cluster>();

            _transparent = new Cluster();
            _semiTransparent = new Cluster();
            _opaque = new Cluster();


        }

        /// <summary>
        /// Extended from base to add initial clusters to the queue
        /// </summary>
        /// <param name="sourceData">The source data</param>
        /// <param name="width">The width in pixels of the image</param>
        /// <param name="height">The height in pixels of the image</param>
        protected override void FirstPass(BitmapData sourceData, int width, int height)
        {
            base.FirstPass(sourceData, width, height);

            // opaque and semitransparent colors are in separate clusters so they will not mingle
            if (_opaque.Pixels.Count > 0) _splittables.Enqueue(_opaque);
            if (_semiTransparent.Pixels.Count > 0) _splittables.Enqueue(_semiTransparent);
            
            // fully transparent pixels are all in the same cluster so there will not be any 
            // duplicate transparent entry in the palette
            if (_transparent.Pixels.Count > 0) _nonsplittables.Add(_transparent);
        }

        /// <summary>
        /// Process the pixel in the first pass of the algorithm
        /// </summary>
        /// <param name="pixel">The pixel to quantize</param>
        protected override void InitialQuantizePixel(Color32* pixel)
        {
            // add pixels to different clusters based on their alpha value
            switch(pixel->Alpha)
            {
                case 0:
                    _transparent.AddPixel(pixel);
                    break;
                case 255:
                    _opaque.AddPixel(pixel);
                    break;
                default:
                    _semiTransparent.AddPixel(pixel);
                    break;
            }        
        }

        /// <summary>
        /// Execute a second pass through the bitmap
        /// </summary>
        /// <remarks>
        /// Perform direct convertion from 32bpp color memory address space to 8bpp indexed space
        /// from cluster to cluster instead of in the original iterative order.
        /// </remarks>
        /// <param name="sourceData">The source bitmap, locked into memory</param>
        /// <param name="output">The output bitmap</param>
        /// <param name="width">The width in pixels of the image</param>
        /// <param name="height">The height in pixels of the image</param>
        /// <param name="bounds">The bounding rectangle</param>
        protected override void SecondPass(BitmapData sourceData, Bitmap output, int width, int height, Rectangle bounds)
        {
            BitmapData outputData = null;

            try
            {
                // Lock the output bitmap into memory
                outputData = output.LockBits(bounds, ImageLockMode.WriteOnly, PixelFormat.Format8bppIndexed);

                byte* pSrc = (byte*)sourceData.Scan0.ToPointer();
                byte* pDst = (byte*)outputData.Scan0.ToPointer();
                int srcStride = sourceData.Stride;
                int dstStride = outputData.Stride;

                int index = 0;
                foreach (Cluster node in _nonsplittables)
                {
                    foreach (IntPtr var in node.Pixels)
                    {
                        // convert from source memory address to destination memory address
                        byte* pSrcPixel = (byte*)var;
                        int srcRow = (int)((pSrcPixel - pSrc) / srcStride);
                        int srcCol = (int)((pSrcPixel - pSrc) % srcStride);
                        byte* pDstIndex = (byte*)(pDst + srcRow * dstStride + srcCol / 4);

                        // assign color palette index
                        *pDstIndex = (byte)index;
                    }
                    index++;
                }
            }
            finally
            {
                // Ensure that I unlock the output bits
                output.UnlockBits(outputData);
            }  
        }

        /// <summary>
        /// Second pass is overridden in such a way that QuantizePixel is no longer used.
        /// </summary>
        /// <param name="pixel">not used</param>
        /// <returns></returns>
        protected override byte QuantizePixel(Color32* pixel)
        {
            throw new Exception("This should not be called");
        }

        /// <summary>
        /// Divide the cluster into two along the longest color dimension.
        /// </summary>
        protected void Split()
        {
            Cluster cluster;
            Color32 medColor;
            IComparer<IntPtr> comparer;

            // get a splittable cluster from the queue
            do
            {
                if (_splittables.Count > 0)
                    cluster = _splittables.Dequeue();
                else
                    return;

                // test if this cluster is splittable
                comparer = cluster.SplitComparer(out medColor);
                if (comparer == null)
                {
                    // move nonsplittable cluster to its list
                    _nonsplittables.Add(cluster);
                    continue;
                }
                else
                {
                    // splittable cluster found
                    break;
                }

            } while (true);


            IntPtr median = (IntPtr)(&medColor);
            Cluster left = new Cluster();
            Cluster right = new Cluster();

            // split cluster into two down the median point
            foreach (IntPtr pPixel in cluster.Pixels)
            {
                if (comparer.Compare(pPixel, median) > 0)
                    left.AddPixel((Color32*)pPixel);
                else
                    right.AddPixel((Color32*)pPixel);
            }

            if (left.Pixels.Count > 0) _splittables.Enqueue(left);
            if (right.Pixels.Count > 0) _splittables.Enqueue(right);
        }

        /// <summary>
        /// Generate color palette through median cut.
        /// </summary>
        /// <param name="original">palette to fill</param>
        /// <returns>new palette</returns>
        protected override ColorPalette GetPalette(ColorPalette original)
        {
            // continue splitting clusters until all clusters are nonsplittable or reaching desired
            // number of clusters            
            while (_splittables.Count + _nonsplittables.Count < _maxColors && _splittables.Count > 0)
                Split();

            // move all clusters into the same list
            while (_splittables.Count > 0)
                _nonsplittables.Add(_splittables.Dequeue());

            // each cluster's average color become a color palette entry
            int index = 0;
            foreach (Cluster cluster in _nonsplittables)
            {
                Color c = cluster.AverageColor();
                original.Entries[index] = c;
                index++;
            }

            return original;
        }

        /// <summary>
        /// Maximum number of colors for the palette, actual number might be lower
        /// </summary>
        private int _maxColors;

        /// <summary>
        /// A queue of clusters that has potential to be further divided
        /// </summary>
        private Queue<Cluster> _splittables;
        
        /// <summary>
        /// A list of clusters that cannot be further divided
        /// </summary>
        private List<Cluster> _nonsplittables;
        
        /// <summary>
        /// Stores all pixels with alpha of 0
        /// </summary>
        private Cluster _transparent;

        /// <summary>
        /// An initial cluster for nonopaque pixels
        /// </summary>
        private Cluster _semiTransparent;
        
        /// <summary>
        /// An initial cluster for opaque pixels
        /// </summary>
        private Cluster _opaque;

        /// <summary>
        /// A collection of pixels that have potentially similar color values.
        /// </summary>
        protected class Cluster
        {
            /// <summary>
            /// Constructs a cluster
            /// </summary>
            public Cluster()
            {
                _pixels = new List<IntPtr>();
                _minAlpha = _minRed = _minGreen = _minBlue = 255;
                _maxAlpha = _maxRed = _maxGreen = _maxBlue = 0;
            }

            /// <summary>
            /// Add address to pixel to the cluster
            /// </summary>
            /// <param name="pixel">pointer to pixel</param>
            public void AddPixel(Color32* pixel)
            {
                _pixels.Add((IntPtr)pixel);

                // scale RGB values by alpha so during quantization, color become less significant to 
                // preserve as pixel gets more transparent
                int postRed = pixel->Red * pixel->Alpha / 255;
                int postGreen = pixel->Green * pixel->Alpha / 255;
                int postBlue = pixel->Blue * pixel->Alpha / 255;

                // stores min and max PARGB values
                if (pixel->Alpha < _minAlpha) _minAlpha = pixel->Alpha;
                if (pixel->Alpha > _maxAlpha) _maxAlpha = pixel->Alpha;

                if (postRed < _minRed) _minRed = postRed;
                if (postRed > _maxRed) _maxRed = postRed;

                if (postGreen < _minGreen) _minGreen = postGreen;
                if (postGreen > _maxGreen) _maxGreen = postGreen;

                if (postBlue < _minBlue) _minBlue = postBlue;
                if (postBlue > _maxBlue) _maxBlue = postBlue;                 
            }

            /// <summary>
            /// Picks the longest dimention and a median point along which the cluster will be cut into two
            /// in argb color space.
            /// </summary>
            /// <param name="median">returns the median point of the cluster</param>
            /// <returns>a comparer that matches the choosen dimention, or null if cluster can't be split</returns>
            public IComparer<IntPtr> SplitComparer(out Color32 median)
            {
                // use average to approximate median
                median = new Color32();
                Color avgColor = AverageColor();
                median.Alpha = avgColor.A;
                median.Red = avgColor.R;
                median.Green = avgColor.G;
                median.Blue = avgColor.B;

                float diff, maxDiff;
                IComparer<IntPtr> comparer;

                // get length in alpha axis
                maxDiff = _maxAlpha - _minAlpha;
                comparer = new SplitByAlpha();
                
                // compare to length in red axis
                diff = (_maxRed - _minRed);
                if (diff > maxDiff)
                {
                    maxDiff = diff;
                    comparer = new SplitByRed();
                }

                // compare to length in green axis
                diff = (_maxGreen - _minGreen);
                if (diff > maxDiff)
                {
                    maxDiff = diff;
                    comparer = new SplitByGreen();
                }

                // compare length in blue axis
                diff = (_maxBlue - _minBlue);
                if (diff > maxDiff)
                {
                    maxDiff = diff;
                    comparer = new SplitByBlue();
                }
                
                // if all dimensions are 0, the cluster has only 1 color and not splittable
                return maxDiff > 0 ? comparer : null;
            }

            /// <summary>
            /// Computes average color among all pixels in the cluster
            /// </summary>
            /// <returns>returns the average color</returns>
            public Color AverageColor()
            {
                int alpha = 0;
                int red = 0;
                int green = 0;
                int blue = 0;
                int count = _pixels.Count;

                checked
                {
                    for (int i = 0; i < count; i++)
                    {
                        Color32* pColor = (Color32*)_pixels[i];
                        alpha += pColor->Alpha;
                        red += pColor->Red;
                        green += pColor->Green;
                        blue += pColor->Blue;
                    }
                }

                return Color.FromArgb(alpha / count, red / count, green / count, blue / count);
            }

            private List<IntPtr> _pixels;
            public List<IntPtr> Pixels
            {
                get { return _pixels; }
            }

            private int _minAlpha, _minRed, _minGreen, _minBlue;
            private int _maxAlpha, _maxRed, _maxGreen, _maxBlue;
        }

        #region Comparers used for KD Tree splitting

        private class SplitByAlpha : IComparer<IntPtr>
        {
            public int Compare(IntPtr x, IntPtr y)
            {
                Color32* c1 = (Color32*)x;
                Color32* c2 = (Color32*)y;
                if (c1->Alpha > c2->Alpha)
                    return 1;
                if (c1->Alpha < c2->Alpha)
                    return -1;
                else
                    return 0;
            }
        }

        private class SplitByRed : IComparer<IntPtr>
        {
            public int Compare(IntPtr x, IntPtr y)
            {
                Color32* c1 = (Color32*)x;
                Color32* c2 = (Color32*)y;
                if (c1->Red > c2->Red)
                    return 1;
                if (c1->Red < c2->Red)
                    return -1;
                else
                    return 0;
            }
        }

        private class SplitByGreen : IComparer<IntPtr>
        {
            public int Compare(IntPtr x, IntPtr y)
            {
                Color32* c1 = (Color32*)x;
                Color32* c2 = (Color32*)y;
                if (c1->Green > c2->Green)
                    return 1;
                if (c1->Green < c2->Green)
                    return -1;
                else
                    return 0;
            }
        }

        private class SplitByBlue : IComparer<IntPtr>
        {
            public int Compare(IntPtr x, IntPtr y)
            {
                Color32* c1 = (Color32*)x;
                Color32* c2 = (Color32*)y;
                if (c1->Blue > c2->Blue)
                    return 1;
                if (c1->Blue < c2->Blue)
                    return -1;
                else
                    return 0;
            }
        }

        #endregion

    }
}
